{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install silentcipher --force-reinstall --no-deps\n",
    "%pip install pydub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import silentcipher\n",
    "import librosa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device='cuda'\n",
    "\n",
    "# You will need to upload the models on colab https://github.com/sony/silentcipher/releases/tag/release\n",
    "\n",
    "model = silentcipher.get_model(\n",
    "    model_type='44.1k',\n",
    "    device=device\n",
    ")\n",
    "\n",
    "\"\"\"\n",
    "For 16khz model\n",
    "\n",
    "model = silentcipher.get_model(\n",
    "    model_type='16k',\n",
    "    device=device\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(y, sr, filename, encoded_filename):\n",
    "    \n",
    "    # The message should be in the form of five 8-bit characters, giving a total message capacity of 40 bits\n",
    "\n",
    "    # Single Channel encoding example\n",
    "    # If you want to really speedup the watermarking process then set disable_checks=True, but beware as this may cause unexpected results\n",
    "    # Set calc_sdr=False to not calculate the SDR of the encoded audio\n",
    "    encoded, sdr = model.encode_wav(y, sr, [123, 234, 111, 222, 11], message_sdr=None, calc_sdr=True, disable_checks=False)\n",
    "\n",
    "    # You can specify the message SDR (in dB) as a float along with the encode_wav function. But this may result in unexpected detection accuracy\n",
    "    # encoded, sdr = model.encode_wav(y, sr, [123, 234, 111, 222, 11], message_sdr=47)\n",
    "\n",
    "    # For multi-channel audio, you can use the following function\n",
    "    # Here y is a 2 channel audio with shape [num_frames, num_channels] and you can specify the message for each channel\n",
    "    # encoded, sdr = model.encode_wav(y, sr, [[123, 234, 111, 222, 11], [132, 214, 121, 122, 211]])\n",
    "\n",
    "    # Single Channel decoding example\n",
    "    result = model.decode_wav(encoded, sr, phase_shift_decoding=False)\n",
    "\n",
    "    if type(result) is list:\n",
    "        for result_i in result:\n",
    "            assert result_i['status']\n",
    "            assert result_i['messages'][0] == [123, 234, 111, 222, 11], result_i['messages'][0]\n",
    "            assert result_i['confidences'][0] == 0.9746031761169434, result_i['confidences'][0]\n",
    "    else:\n",
    "        assert result['status']\n",
    "        assert result['messages'][0] == [123, 234, 111, 222, 11], result['messages'][0]\n",
    "        assert result['confidences'][0] == 0.9746031761169434, result['confidences'][0]\n",
    "\n",
    "    # When decoding multi-channel audio, The result would be a list of dictionaries with the status, message and confidence for each channel\n",
    "\n",
    "    # Encode from filename\n",
    "\n",
    "    model.encode(filename, encoded_filename, [123, 234, 111, 222, 11], message_sdr=None, calc_sdr=True, disable_checks=False)\n",
    "    result = model.decode(encoded_filename, phase_shift_decoding=False)\n",
    "\n",
    "    if type(result) is list:\n",
    "        for result_i in result:\n",
    "            assert result_i['status']\n",
    "            assert result_i['messages'][0] == [123, 234, 111, 222, 11], result_i['messages'][0]\n",
    "            assert result_i['confidences'][0] == 0.9746031761169434, result_i['confidences'][0]\n",
    "    else:\n",
    "        assert result['status']\n",
    "        assert result['messages'][0] == [123, 234, 111, 222, 11], result['messages'][0]\n",
    "        assert result['confidences'][0] == 0.9746031761169434, result['confidences'][0]\n",
    "\n",
    "    # When decoding multi-channel audio, The result would be a list of dictionaries with the status, message and confidence for each channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You will need to upload the test.wav and test_multichannel.wav from the examples/colab directory\n",
    "\n",
    "y, sr = librosa.load('test.wav', sr=None)\n",
    "test(y, sr, 'test.wav', 'encoded_test.wav')\n",
    "y, sr = librosa.load('test_multichannel.wav', sr=None)\n",
    "test(y, sr, 'test_multichannel.wav', 'encoded_test_multichannel.wav')\n",
    "\n",
    "print('Test Successful')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
